# -*- coding: utf-8 -*-
"""
    Project name：code_potentialflow
# -------------------------------
    File name：transformations.py.py
    Created on：2025/6/7 15:36
    Author：(input)
    Description:区域变换类
"""
import numpy as np
import sympy as sp
from scipy.special import ellipk, ellipj

# ----------------------
# Base Class for Transformations
# ----------------------
class BaseMapping:
    def forward(self, z):
        raise NotImplementedError

    def inverse(self, w):
        raise NotImplementedError


# ----------------------
# Möbius Transform Class
# ----------------------
class MobiusTransform(BaseMapping):
    def __init__(self, z_vals, w_vals):
        self.z_vals = z_vals
        self.w_vals = w_vals
        self.z = sp.Symbol('z', complex=True)
        self.a, self.b, self.d, self.beta = sp.symbols('a b d beta', real=True)
        self.f_expr = (self.a * self.z + self.b) / (self.z + self.d)
        self.solution = None
        self._solve()

    def _solve(self):
        equations = [
            sp.Eq(self.f_expr.subs(self.z, z), w)
            for z, w in zip(self.z_vals, self.w_vals)
        ]
        solutions = sp.solve(equations, (self.a, self.b, self.d, self.beta), dict=True)

        for sol in solutions:
            beta_val = sol[self.beta]
            if 0 < sp.re(beta_val) < 1:
                self.solution = sol
                break

        if self.solution is None:
            raise ValueError("No valid Möbius solution with 0 < beta < 1")

        self.f_final = self.f_expr.subs(self.solution)

    def forward(self, z_val):
        z = sp.sympify(z_val)
        return complex(sp.N(self.f_final.subs(self.z, z)))

    def forward_array(self, Z):
        return np.vectorize(self.forward)(Z)


# ----------------------
# Bilinear Transform (2D quad map)
# ----------------------
def bilinear_map(xi, eta, points):
    """
    Bilinear map from (xi, eta) ∈ [0,1]² to physical quad
    points: list of 4 points [A1, A2, A3, A4]
    """
    N0 = xi * eta        # A1
    N1 = (1 - xi) * eta  # A2
    N2 = (1 - xi) * (1 - eta)  # A3
    N3 = xi * (1 - eta)  # A4
    shape = np.array([N0, N1, N2, N3])
    points = np.array(points)
    x = np.dot(shape, points[:, 0])
    y = np.dot(shape, points[:, 1])
    return x + 1j * y


# ----------------------
# SC Transform Placeholder
# ----------------------
from core.sc_rect_solver import ParameterSolver

class SchwarzChristoffelTransform:
    def __init__(self, polygon_vertices=None, target_domain='half_plane'):
        self.vertices = polygon_vertices
        self.target_domain = target_domain

    def forward(self, z):
        raise NotImplementedError

    def inverse(self, w):
        raise NotImplementedError


class SCRectangleTransform(SchwarzChristoffelTransform):
    """
    将矩形区域（底边中心对称）映射到上半平面的 SC 变换
    """
    def __init__(self, width, height, characteristic=1.0, dimensional=True):
        """
        初始化参数求解器，输入width、height，特征长度characteristic，有无量纲dimensional
        width: 长度参数
        height: 高度参数
        characteristic：特征长度
        dimensional: 输入是否为有量纲量
        """
        super().__init__(polygon_vertices=None, target_domain='half_plane')

        self.width = width
        self.height = height
        self.characteristic = characteristic
        self.dimensional = dimensional

        # 求解参数（C, k, k'）
        self.solver = ParameterSolver(width, height,
                                      characteristic=characteristic,
                                      dimensional=dimensional)
        self.C, self.k, self.kp = self.solver.solve()

        # 计算 K(k), K'(k)
        self.Kk = ellipk(self.k ** 2)
        self.Kkp = ellipk(self.kp ** 2)

    def forward(self, z):
        u = self.C * z
        # 确保参数是浮点数
        u_float = np.float64(u)
        m = np.float64(self.k ** 2)
        sn, _, _, _ = ellipj(u_float, m)
        return sn

    def forward_array(self, Z):
        """对数组中的每个元素单独处理"""
        Z_array = np.asarray(Z)
        result = np.empty_like(Z_array, dtype=np.float64)

        for idx in np.ndindex(Z_array.shape):
            u = self.C * Z_array[idx]
            m = np.float64(self.k ** 2)
            sn, _, _, _ = ellipj(u, m)
            result[idx] = sn

        return result

    def get_mapping_info(self):
        return {
            'C': self.C,
            'k': self.k,
            "k'": self.kp,
            'K(k)': self.Kk,
            "K(k')": self.Kkp
        }
